import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import akshare as ak
import jieba
from textblob import TextBlob
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.ensemble import RandomForestClassifier
import time
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')

class MultiDimensionStockAnalyzer:
    def __init__(self):
        """初始化多维度股票分析器"""
        self.stock_data = {}  # 股票数据
        self.news_data = {}  # 新闻数据
        self.sentiment_model = None  # 情感分析模型
        self.vectorizer = None  # 文本向量化器
        
    def fetch_stock_data(self, stock_code, start_date, end_date):
        """
        获取股票历史数据
        :param stock_code: 股票代码，如 'sh000001'
        :param start_date: 开始日期，格式 '20200101'
        :param end_date: 结束日期，格式 '20230101'
        """
        try:
            # 使用akshare获取股票日线数据
            stock_df = ak.stock_zh_a_hist(
                symbol=stock_code[2:], 
                period="daily", 
                start_date=start_date, 
                end_date=end_date, 
                adjust="qfq"
            )
            
            # 重命名列名以便于处理
            # stock_df.columns = ['日期', '开盘', '收盘', '最高', '最低', '成交量', '成交额', '振幅', '涨跌幅', '涨跌额', '换手率']
            
            # 转换日期格式
            stock_df['日期'] = pd.to_datetime(stock_df['日期'])
            
            self.stock_data[stock_code] = stock_df
            print(f"成功获取{stock_code}的股票数据，共{len(stock_df)}行记录")
            return stock_df
        except Exception as e:
            print(f"获取股票数据失败: {e}")
            return pd.DataFrame()
    
    def fetch_news(self, stock_code, start_date, end_date, keyword=None):
        """
        获取指定股票的相关新闻
        :param stock_code: 股票代码
        :param start_date: 开始日期
        :param end_date: 结束日期
        :param keyword: 关键词，默认为股票名称
        """
        try:
            # 使用akshare获取财经新闻
            # 注：实际应用中可能需要结合多个API或网页爬虫获取更全面的新闻
            news_list = []
            
            # 模拟获取新闻数据
            for i in range(5):  # 获取5条示例新闻
                date = (datetime.strptime(end_date, '%Y%m%d') - timedelta(days=i)).strftime('%Y-%m-%d')
                news_item = {
                    'date': date,
                    'title': f"{stock_code}公司发布新产品，市场前景广阔",
                    'content': f"{date}日，{stock_code}公司宣布推出一款创新产品，预计将在未来季度带来显著收入增长。"
                             f"分析师认为该产品将提升公司市场份额。同时，公司计划扩大生产规模，以满足市场需求。",
                    'source': '财经新闻网'
                }
                news_list.append(news_item)
            
            news_df = pd.DataFrame(news_list)
            self.news_data[stock_code] = news_df
            return news_df
        except Exception as e:
            print(f"获取新闻数据失败: {e}")
            return pd.DataFrame()
    
    def preprocess_news(self, news_df):
        """
        预处理新闻文本
        :param news_df: 新闻数据框
        :return: 处理后的新闻数据框
        """
        if news_df.empty:
            return news_df
        
        # 分词处理（中文）
        def chinese_tokenize(text):
            return jieba.cut(text)
        
        # 计算新闻情感得分
        news_df['sentiment'] = news_df['content'].apply(lambda x: TextBlob(x).sentiment.polarity)
        
        # 分词并创建词袋
        news_df['tokens'] = news_df['content'].apply(chinese_tokenize)
        
        return news_df
    
    def train_sentiment_model(self, news_df=None):
        """
        训练新闻情感分析模型
        :param news_df: 训练数据，默认为已获取的新闻
        """
        if news_df is None:
            # 使用所有已获取的新闻数据
            all_news = pd.concat(self.news_data.values(), ignore_index=True) if self.news_data else pd.DataFrame()
            if all_news.empty:
                print("没有足够的新闻数据用于训练情感模型")
                return False
        else:
            all_news = news_df
        
        # 假设我们有带标签的训练数据（实际应用中需要真实标注数据）
        # 这里使用模拟标签进行演示
        all_news['label'] = np.where(all_news['sentiment'] > 0, 1, 
                                    np.where(all_news['sentiment'] < 0, -1, 0))
        
        # 文本向量化
        self.vectorizer = TfidfVectorizer(tokenizer=jieba.cut)
        X = self.vectorizer.fit_transform(all_news['content'])
        y = all_news['label']
        
        # 训练随机森林分类器
        self.sentiment_model = RandomForestClassifier(n_estimators=100, random_state=42)
        self.sentiment_model.fit(X, y)
        
        print("情感分析模型训练完成")
        return True
    
    def analyze_news_sentiment(self, news_df):
        """
        分析新闻情感
        :param news_df: 新闻数据框
        :return: 带情感分析结果的数据框
        """
        if news_df.empty or self.sentiment_model is None or self.vectorizer is None:
            print("无法进行情感分析：模型或数据不足")
            return news_df
        
        # 向量化文本
        X = self.vectorizer.transform(news_df['content'])
        
        # 预测情感
        news_df['predicted_sentiment'] = self.sentiment_model.predict(X)
        news_df['sentiment_score'] = self.sentiment_model.predict_proba(X).max(axis=1)
        
        return news_df
    
    def analyze_market_dimension(self, stock_code):
        """
        分析市场维度（量价关系）
        :param stock_code: 股票代码
        :return: 分析结果
        """
        if stock_code not in self.stock_data:
            print(f"没有{stock_code}的股票数据")
            return {}
        
        df = self.stock_data[stock_code].copy()
        
        # 计算技术指标
        df['MA5'] = df['收盘'].rolling(window=5).mean()
        df['MA10'] = df['收盘'].rolling(window=10).mean()
        df['MA20'] = df['收盘'].rolling(window=20).mean()
        
        # 计算成交量变化
        df['volume_change'] = df['成交量'].pct_change()
        
        # 计算价格变化
        df['price_change'] = df['收盘'].pct_change()
        
        # 分析量价关系
        df['price_volume_corr'] = df['price_change'].rolling(window=20).corr(df['volume_change'])
        
        # 识别异常交易
        high_volume_days = df[df['volume_change'] > df['volume_change'].mean() + 2 * df['volume_change'].std()]
        
        analysis = {
            'trend': '上升' if df['收盘'].iloc[-1] > df['MA20'].iloc[-1] else '下降',
            'volume_trend': '放量' if df['成交量'].iloc[-1] > df['成交量'].rolling(window=5).mean().iloc[-1] else '缩量',
            'price_volume_corr': df['price_volume_corr'].iloc[-1],
            'high_volume_days': high_volume_days[['日期', '收盘', '成交量', 'price_change', 'volume_change']]
        }
        
        return analysis
    
    def analyze_news_dimension(self, stock_code):
        """
        分析新闻维度
        :param stock_code: 股票代码
        :return: 分析结果
        """
        if stock_code not in self.news_data:
            print(f"没有{stock_code}的新闻数据")
            return {}
        
        news_df = self.news_data[stock_code].copy()
        
        # 预处理新闻
        news_df = self.preprocess_news(news_df)
        
        # 分析情感
        if self.sentiment_model and self.vectorizer:
            news_df = self.analyze_news_sentiment(news_df)
        else:
            # 使用简单情感分析
            news_df['predicted_sentiment'] = np.where(news_df['sentiment'] > 0, 1, 
                                                     np.where(news_df['sentiment'] < 0, -1, 0))
        
        # 按日期分组统计情感
        sentiment_by_date = news_df.groupby('date')['predicted_sentiment'].mean().reset_index()
        
        # 识别重要新闻
        important_news = news_df[news_df['sentiment'].abs() > 0.5]
        
        analysis = {
            'avg_sentiment': sentiment_by_date['predicted_sentiment'].mean(),
            'sentiment_trend': sentiment_by_date,
            'positive_news': news_df[news_df['predicted_sentiment'] > 0],
            'negative_news': news_df[news_df['predicted_sentiment'] < 0],
            'important_news': important_news
        }
        
        return analysis
    
    def analyze_time_dimension(self, stock_code, event_dates=None):
        """
        分析时间维度（事件与价格的时间关系）
        :param stock_code: 股票代码
        :param event_dates: 重要事件日期列表
        :return: 分析结果
        """
        if stock_code not in self.stock_data:
            print(f"没有{stock_code}的股票数据")
            return {}
        
        df = self.stock_data[stock_code].copy()
        
        if event_dates is None and stock_code in self.news_data:
            # 使用新闻日期作为事件日期
            event_dates = self.news_data[stock_code]['date'].unique().tolist()
        
        if not event_dates:
            print("没有指定事件日期")
            return {}
        
        # 分析事件前后的价格变化
        event_impacts = []
        
        for date in event_dates:
            try:
                # 查找最接近的交易日
                closest_trading_day = df.iloc[(df['日期'] - pd.to_datetime(date)).abs().argsort()[0]]['日期']
                idx = df[df['日期'] == closest_trading_day].index[0]
                
                # 计算事件前后的价格变化
                if idx >= 5 and idx < len(df) - 5:
                    prev_price = df.iloc[idx-5]['收盘']
                    event_price = df.iloc[idx]['收盘']
                    post_price = df.iloc[idx+5]['收盘']
                    
                    impact = {
                        'event_date': date,
                        'trading_date': closest_trading_day,
                        'price_before': prev_price,
                        'price_at_event': event_price,
                        'price_after': post_price,
                        'change_before': (event_price - prev_price) / prev_price * 100,
                        'change_after': (post_price - event_price) / event_price * 100
                    }
                    event_impacts.append(impact)
            except Exception as e:
                print(f"分析事件日期 {date} 时出错: {e}")
        
        return pd.DataFrame(event_impacts) if event_impacts else pd.DataFrame()
    
    def analyze_cross_dimension(self, stock_code):
        """
        分析多维度交叉关系
        :param stock_code: 股票代码
        :return: 交叉分析结果
        """
        market_analysis = self.analyze_market_dimension(stock_code)
        news_analysis = self.analyze_news_dimension(stock_code)
        
        if not market_analysis or not news_analysis:
            print("数据不足，无法进行交叉维度分析")
            return {}
        
        # 合并新闻情感与股价数据
        if 'sentiment_trend' in news_analysis and not news_analysis['sentiment_trend'].empty:
            sentiment_df = news_analysis['sentiment_trend']
            sentiment_df['date'] = pd.to_datetime(sentiment_df['date'])
            
            stock_df = self.stock_data[stock_code].copy()
            stock_df['日期'] = pd.to_datetime(stock_df['日期'])
            
            # 按日期合并
            merged_df = pd.merge_asof(
                stock_df.sort_values('日期'),
                sentiment_df.sort_values('date'),
                left_on='日期',
                right_on='date',
                tolerance=pd.Timedelta(days=1),
                direction='nearest'
            )
            
            # 计算情感与价格变化的相关性
            merged_df['price_change'] = merged_df['收盘'].pct_change()
            sentiment_price_corr = merged_df['predicted_sentiment'].corr(merged_df['price_change'].shift(-1))
            
            # 分析情感领先价格变动的情况
            sentiment_leading_positive = merged_df[
                (merged_df['predicted_sentiment'] > 0) & 
                (merged_df['price_change'].shift(-1) > 0)
            ]
            
            sentiment_leading_negative = merged_df[
                (merged_df['predicted_sentiment'] < 0) & 
                (merged_df['price_change'].shift(-1) < 0)
            ]
            
            cross_analysis = {
                'sentiment_price_corr': sentiment_price_corr,
                'sentiment_leading_positive': len(sentiment_leading_positive) / len(merged_df.dropna()) * 100,
                'sentiment_leading_negative': len(sentiment_leading_negative) / len(merged_df.dropna()) * 100,
                'merged_data': merged_df
            }
            
            return cross_analysis
        
        return {}
    
    def visualize_multi_dimension(self, stock_code):
        """
        可视化多维度分析结果
        :param stock_code: 股票代码
        """
        if stock_code not in self.stock_data:
            print(f"没有{stock_code}的股票数据")
            return
        
        plt.figure(figsize=(16, 12))
        
        # 1. 绘制股票价格与均线
        plt.subplot(2, 2, 1)
        df = self.stock_data[stock_code].copy()
        plt.plot(df['日期'], df['收盘'], label='收盘价')
        plt.plot(df['日期'], df['收盘'].rolling(window=5).mean(), label='5日均线')
        plt.plot(df['日期'], df['收盘'].rolling(window=20).mean(), label='20日均线')
        plt.title(f'{stock_code} 股价与均线走势')
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True)
        
        # 2. 绘制成交量与价格变化
        plt.subplot(2, 2, 2)
        plt.bar(df['日期'], df['成交量'], alpha=0.5, label='成交量')
        plt.twinx()
        plt.plot(df['日期'], df['收盘'].pct_change() * 100, 'r-', label='价格变化率(%)')
        plt.title(f'{stock_code} 成交量与价格变化')
        plt.xlabel('日期')
        plt.legend()
        plt.grid(True)
        
        # 3. 绘制新闻情感分析
        if stock_code in self.news_data and not self.news_data[stock_code].empty:
            news_df = self.news_data[stock_code].copy()
            news_df = self.preprocess_news(news_df)
            
            if self.sentiment_model and self.vectorizer:
                news_df = self.analyze_news_sentiment(news_df)
            else:
                news_df['predicted_sentiment'] = np.where(news_df['sentiment'] > 0, 1, 
                                                         np.where(news_df['sentiment'] < 0, -1, 0))
            
            sentiment_by_date = news_df.groupby('date')['predicted_sentiment'].mean().reset_index()
            sentiment_by_date['date'] = pd.to_datetime(sentiment_by_date['date'])
            
            plt.subplot(2, 2, 3)
            plt.bar(sentiment_by_date['date'], sentiment_by_date['predicted_sentiment'], color='g')
            plt.axhline(y=0, color='r', linestyle='-')
            plt.title(f'{stock_code} 新闻情感分析')
            plt.xlabel('日期')
            plt.ylabel('情感分数')
            plt.grid(True)
        
        # 4. 绘制多维度交叉分析
        cross_analysis = self.analyze_cross_dimension(stock_code)
        if cross_analysis and 'merged_data' in cross_analysis:
            merged_df = cross_analysis['merged_data'].dropna()
            
            plt.subplot(2, 2, 4)
            plt.scatter(merged_df['predicted_sentiment'], merged_df['price_change'].shift(-1), alpha=0.5)
            plt.axhline(y=0, color='r', linestyle='-')
            plt.axvline(x=0, color='r', linestyle='-')
            plt.title(f'新闻情感与未来价格变化的关系 (相关系数: {cross_analysis["sentiment_price_corr"]:.4f})')
            plt.xlabel('情感分数')
            plt.ylabel('未来价格变化率')
            plt.grid(True)
        
        plt.tight_layout()
        plt.show()


def main():
    # 创建多维度股票分析器
    analyzer = MultiDimensionStockAnalyzer()
    
    # 设置参数
    stock_code = 'sh601318'  # 中国平安
    start_date = '20230101'
    end_date = '20230701'
    
    # 1. 获取股票数据
    stock_data = analyzer.fetch_stock_data(stock_code, start_date, end_date)
    if stock_data.empty:
        return
    
    # 2. 获取新闻数据
    news_data = analyzer.fetch_news(stock_code, start_date, end_date)
    if news_data.empty:
        print("无法获取新闻数据，部分分析将受限")
    
    # 3. 训练情感分析模型
    analyzer.train_sentiment_model()
    
    # 4. 市场维度分析
    market_analysis = analyzer.analyze_market_dimension(stock_code)
    print("\n市场维度分析结果:")
    print(f"趋势: {market_analysis['trend']}")
    print(f"成交量趋势: {market_analysis['volume_trend']}")
    print(f"量价相关性: {market_analysis['price_volume_corr']:.4f}")
    
    # 5. 新闻维度分析
    news_analysis = analyzer.analyze_news_dimension(stock_code)
    print("\n新闻维度分析结果:")
    print(f"平均情感分数: {news_analysis['avg_sentiment']:.4f}")
    print(f"正面新闻数量: {len(news_analysis['positive_news'])}")
    print(f"负面新闻数量: {len(news_analysis['negative_news'])}")
    
    # 6. 时间维度分析
    time_analysis = analyzer.analyze_time_dimension(stock_code)
    if not time_analysis.empty:
        print("\n时间维度分析结果:")
        print(f"分析了{len(time_analysis)}个事件")
        print(f"事件前平均价格变化: {time_analysis['change_before'].mean():.2f}%")
        print(f"事件后平均价格变化: {time_analysis['change_after'].mean():.2f}%")
    
    # 7. 交叉维度分析
    cross_analysis = analyzer.analyze_cross_dimension(stock_code)
    if cross_analysis and 'sentiment_price_corr' in cross_analysis:
        print("\n交叉维度分析结果:")
        print(f"新闻情感与未来价格变化的相关性: {cross_analysis['sentiment_price_corr']:.4f}")
        print(f"正面情感预测上涨的准确率: {cross_analysis['sentiment_leading_positive']:.2f}%")
        print(f"负面情感预测下跌的准确率: {cross_analysis['sentiment_leading_negative']:.2f}%")
    
    # 8. 可视化多维度分析
    analyzer.visualize_multi_dimension(stock_code)


if __name__ == "__main__":
    main()    
